from matplotlib import rcParams
import matplotlib
matplotlib.use('TkAgg')  # 或者 'Agg'
# 设置字体为 SimHei（黑体），确保中文显示正常
rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

import os
import time
import random
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import config  # 确保 config.py 存在，并且包含 OUTPUT_FOLDER_ONE_BLEW

def draw_heatmap(file_path, select_feature, cmap="coolwarm"):
    """
    绘制上三角热图，显示特征之间的相关性。

    参数：
    - file_path: 输入文件路径，支持 Excel 格式。
    - select_feature: 目标列名称列表，用于选择需要分析的特征。
    - cmap (str): 热图的配色方案，默认为 "coolwarm"。

    返回：
    - output_file: 生成的热图文件路径。
    """
    # 生成唯一文件名
    timestamp = int(time.time() * 1000)
    random_number = random.randint(1000, 9999)
    output_filename = f"{timestamp}_{random_number}.png"

    # 确保输出文件夹存在
    output_dir = config.OUTPUT_FOLDER_ONE_BLEW
    os.makedirs(output_dir, exist_ok=True)

    # 输出文件路径
    output_file = os.path.join(output_dir, output_filename)

    # 读取数据
    df = pd.read_excel(file_path)
    df = df[select_feature]  # 选择目标特征列

    # 计算相关矩阵
    correlation_matrix = df.corr()

    # 创建画布
    fig, ax = plt.subplots(figsize=(10, 8))

    # 绘制上三角热图
   # mask = np.tril(np.ones_like(correlation_matrix, dtype=bool), k=-1)
    sns.heatmap(
        correlation_matrix,
       # mask=mask,
        annot=True,
        fmt=".2f",
        cmap=cmap,
        cbar=False,
        ax=ax
    )

    # 调整标签位置
    ax.xaxis.tick_top()
    ax.xaxis.set_label_position('top')
    ax.tick_params(axis='x', which='major', pad=10, length=0)
    plt.xticks(rotation=90)

    ax.yaxis.tick_right()
    ax.yaxis.set_label_position('right')
    ax.tick_params(axis='y', which='major', pad=10, length=0)
    plt.yticks(rotation=0)

    # 调整布局
    plt.tight_layout()
    plt.show()

    # 保存图像
    fig.savefig(output_file)

    print(f"热图已保存至: {output_file}")  # 确保 print 在 return 之前
    return output_file




# def draw_scatter(file_path, select_feature, output_path='./scatter.png', s=1):
#     """
#     绘制下三角散点图，显示特征之间的分布关系，并返回散点数据。
#
#     参数：
#     - file_path: 输入文件路径，支持 Excel 格式。
#     - select_feature: 目标列名称列表，用于选择需要分析的特征。
#     - output_path: 输出图像保存路径，默认为 'scatter.png'。
#     - s (int): 散点的大小，默认为 1。
#
#     返回：
#     - scatter_data: 包含所有散点图数据的字典，键为特征对，值为 (x, y) 坐标数组。
#     """
#     # 1. 读取数据
#     df = pd.read_excel(file_path)
#     df = df[select_feature]  # 选择目标特征列
#
#     # 2. 获取变量数量
#     n = len(df.columns)
#
#     # 3. 设置画布和网格布局
#     fig, axes = plt.subplots(n, n, figsize=(12, 12))
#
#     # 4. 初始化散点数据字典
#     scatter_data = {}
#
#     # 5. 遍历矩阵，绘制下三角散点图
#     for i in range(n):
#         for j in range(n):
#             ax = axes[i, j]  # 获取当前子图
#             if i >= j:  # 下三角部分绘制散点图
#                 x = df.iloc[:, j]
#                 y = df.iloc[:, i]
#                 ax.scatter(x, y, alpha=0.5, s=s, color="blue")  # 绘制散点图
#                 ax.set_xticks([])  # 隐藏x轴刻度
#                 ax.set_yticks([])  # 隐藏y轴刻度
#
#                 # 将散点数据存储到字典中
#                 scatter_data[f"{df.columns[j]} vs {df.columns[i]}"] = (x.values, y.values)
#             else:  # 上三角部分留空
#                 ax.axis('off')
#
#             # 6. 添加列名标签
#             if i == n - 1:  # 最底部的图添加 x 轴标签
#                 ax.set_xlabel(df.columns[j], fontsize=8, rotation=90)
#             if j == 0:  # 最左侧的图添加 y 轴标签
#                 ax.set_ylabel(df.columns[i], fontsize=8, rotation=0, labelpad=40)
#
#     # 7. 调整布局
#     plt.subplots_adjust(wspace=0, hspace=0)  # 消除网格间隙
#
#     # 8. 保存图像
#     fig.savefig(output_path)
#     print(f"散点图已保存至 下三角散点图: {output_path}")
#
#     # 9. 打印散点数据
#     print("散点数据如下 下三角散点图：")
#     for key, (x, y) in scatter_data.items():
#         print(f"{key}:")
#         print(f"x: {x.tolist()}")
#         print(f"y: {y.tolist()}")
#         print()
#
#     return output_path
#



# 绘制上三角热图
# draw_heatmap(
#     file_path='抗压强度.xlsx',  # 数据文件路径
#     select_feature=["水用量（kg/m3）", "水泥ID", "水泥用量（kg/m3）", "粉煤灰用量（kg/m3）", "砂ID", "砂用量（kg/m3）",
#                     "石ID", "石用量（kg/m3）", "减水剂ID", "减水剂掺量（%）", "增效剂ID", "增效剂掺量（%）"
#                     ],  # 目标特征列
#     output_path='heatmap_output.png',  # 输出文件路径
# )

# 绘制下三角散点图
# draw_scatter(
#     file_path='抗压强度.xlsx',  # 数据文件路径
#     select_feature=["水用量（kg/m3）", "水泥ID", "水泥用量（kg/m3）", "粉煤灰用量（kg/m3）", "砂ID", "砂用量（kg/m3）",
#                     "石ID", "石用量（kg/m3）", "减水剂ID", "减水剂掺量（%）", "增效剂ID", "增效剂掺量（%）"
#                     ],  # 目标特征列
#     output_path='scatter_output.png',  # 输出文件路径
# )
